import torch.nn as nn
import torch.nn.functional as F


class ResEnc(nn.Module):
    def __init__(self, in_dementions, out_dementions, kernel_size=1, padding=0, pooling=False):
        super().__init__()
        self.link = (in_dementions != out_dementions)
        self.pooling = pooling
        if self.link:
            self.conv_link = nn.Conv2d(in_dementions, out_dementions, kernel_size=1)
            self.bn_link = nn.BatchNorm2d(out_dementions)

        self.conv1 = nn.Conv2d(in_dementions, out_dementions, kernel_size=kernel_size, padding=padding)
        self.bn1 = nn.BatchNorm2d(out_dementions)
        if pooling:
            self.pool = nn.MaxPool2d(2)

        self.conv2 = nn.Conv2d(out_dementions, out_dementions, kernel_size=1)
        self.bn2 = nn.BatchNorm2d(out_dementions)

    def forward(self, input):
        res = input
        if self.link:
            res = self.conv_link(res)
            res = self.bn_link(res)

        x1 = F.relu(self.bn1(self.conv1(input)))
        if self.pooling:
            x1 = self.pool(x1)
            res = self.pool(res)
        x2 = self.bn2(self.conv2(x1))
        x2 += res
        output = F.relu(x2)

        return output


class Encoder(nn.Module):
    def __init__(self, filters, in_dementions, bias=True, name='encoder'):
        super().__init__()
        self.bn0 = nn.BatchNorm2d(in_dementions)
        self.conv1 = nn.Conv2d(in_dementions, filters[0], kernel_size=1)
        self.bn1 = nn.BatchNorm2d(filters[0])
        self.res1 = ResEnc(filters[0], filters[1], kernel_size=3, padding=1)
        self.res2 = ResEnc(filters[1], filters[2], kernel_size=3, padding=1, pooling=True)
        self.res3 = ResEnc(filters[2], filters[3], kernel_size=3, padding=1, pooling=True)
        self.res4 = ResEnc(filters[3], filters[4], pooling=True)
        # self.res1 = ResEnc(filters[0],filters[1],kernel_size=1)
        # self.res2 = ResEnc(filters[1],filters[2], kernel_size=3, padding=1, pooling=True)
        # self.res3 = ResEnc(filters[2], filters[3], kernel_size=1, pooling=True)
        # self.res4 = ResEnc(filters[3], filters[4], pooling=True)

    def forward(self, input):
        layer0 = F.relu(self.bn0(input))
        layer1 = F.relu(self.bn1(self.conv1(layer0)))

        # attention: we can double the resblock each layer
        # layer2 size 128
        layer2 = self.res1(layer1)
        # layer3 size 64
        layer3 = self.res2(layer2)
        # layer4 size 32
        layer4 = self.res3(layer3)
        # layer5 size 16
        layer5 = self.res4(layer4)

        return [input, layer1, layer2, layer3, layer4, layer5]
